package io.growing.adaptor.common.async;

import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Supplier;

/**
 * @ClassName 与PriorityDistributedBlockingQueue配套的线程池 @Description TODO @Author wangjiu @Date
 * 2021/12/8 7:31 下午
 */
public class PriorityDistributedThreadPoolExecutor extends ThreadPoolExecutor {

    private volatile RejectedExecutionHandler handler;

    private final PriorityDistributedBlockingQueue workQueue;

    public PriorityDistributedThreadPoolExecutor(
            int corePoolSize,
            int maximumPoolSize,
            long keepAliveTime,
            TimeUnit unit,
            Map<Integer, BlockingQueue<Runnable>> queueMap,
            ThreadFactory threadFactory) {
        this(
                corePoolSize,
                maximumPoolSize,
                keepAliveTime,
                unit,
                new PriorityDistributedBlockingQueue(queueMap),
                threadFactory,
                new CallerRunsPolicy());
    }

    public PriorityDistributedThreadPoolExecutor(
            int corePoolSize,
            int maximumPoolSize,
            long keepAliveTime,
            TimeUnit unit,
            PriorityDistributedBlockingQueue workQueue,
            ThreadFactory threadFactory) {
        this(
                corePoolSize,
                maximumPoolSize,
                keepAliveTime,
                unit,
                workQueue,
                threadFactory,
                new CallerRunsPolicy());
    }

    private PriorityDistributedThreadPoolExecutor(
            int corePoolSize,
            int maximumPoolSize,
            long keepAliveTime,
            TimeUnit unit,
            PriorityDistributedBlockingQueue workQueue,
            ThreadFactory threadFactory,
            RejectedExecutionHandler handler) {
        super(
                corePoolSize,
                maximumPoolSize,
                keepAliveTime,
                unit,
                workQueue,
                threadFactory,
                handler);
        this.handler = handler;
        this.workQueue = workQueue;
    }

    @Override
    public void execute(Runnable command) {
        if (!(command instanceof PriorityRunnable)) {
            // 默认是最低优先级，这样就不会影响到其他带优先级的任务执行
            command = new PriorityRunnable(command, Integer.MAX_VALUE);
        }
        super.execute(command);
    }

    public void execute(Runnable command, int priority) {
        if (command == null) throw new NullPointerException();
        if (!(command instanceof PriorityRunnable)) {
            command = new PriorityRunnable(command, priority);
        }
        if (super.getPoolSize() < super.getMaximumPoolSize()
                || priority <= workQueue.getMinPriority()) {
            // 刚开始时，线程数没有达到最大，这时候认为线程资源不紧张，没必要按优先级区别对待任务，
            execute(command);
        } else if (!workQueue.offer(command, priority)) {
            handler.rejectedExecution(command, this);
        }
    }

    public <U> CompletableFuture<U> execute(Supplier<U> supplier, int priority) {
        return FutureUtils.supplyAsync(supplier, this, priority);
    }
}
